﻿//
// An implementation of the Spread client protocol for FSharp.
//
// For more details on the Spread Toolkit, see http://www.spread.org.
//
// (c) 2009 Nick Gunn. http://iodyner.spaces.live.com
//


#nowarn "40"
#nowarn "62"

open System
open System.IO
open System.Net
open System.Net.Sockets

open Microsoft.FSharp.Control
open Microsoft.FSharp.Control.CommonExtensions

// Public types
[<Flags>]
type MessageFlags = 
    | None       = 0x00000000
    | Unreliable = 0x00000001
    | Reliable = 0x00000002
    | Fifo = 0x00000004
    | Causal = 0x00000008
    | Agreed = 0x00000010
    | Safe = 0x00000020
    | Regular = 0x0000003f

    | NotRegular = 0x00ffff00

    | SelfDiscard = 0x00000040

    | RegMembership = 0x00001000
    | Transition = 0x00002000
    | CausedByJoin = 0x00000100
    | CausedByLeave = 0x00000200
    | CausedByDisconnect = 0x00000400
    | CausedByNetwork = 0x00000800
    | Membership = 0x00003f00

    | Join = 0x00010000
    | Leave = 0x00020000
    | Kill = 0x00040000
    | Groups = 0x00080000

    | Reject = 0x00100000

type GroupId = {
    ProcessId : int
    Timestamp : int
    Index : int }

type ChangeSet = {
    Gid : GroupId
    IndexOfLocalGroup : int
    GroupMembers : string array array }

type RegularOutbound = {
    Groups : string array 
    ServiceType : MessageFlags 
    MessageType : uint16 
    Data : byte array }

type RegularInbound = {
    Sender : string
    Groups : string seq 
    MessageType : uint16 
    Data : byte array 
    SameEndian : bool }

type TransitionMessage = {
    Group : string
    Gid : GroupId }

type MembershipReason = 
    | Joined of string
    | Left of string
    | Disconnected of string
    | Partition of ChangeSet
    
type MembershipMessage = {
    Group   : string
    Members : string seq 
    Reason  : MembershipReason }

// An outbound message    
type SpreadMessage = 
    | Data of RegularOutbound
    | Join of string
    | Leave of string
    | Disconnect
    | Kill

// An inbound event    
type SpreadEvent = 
    | Connection of string * (SpreadMessage -> unit)
    | Disconnection of exn option
    | Message of RegularInbound
    | Transition of TransitionMessage
    | MembershipChange of MembershipMessage
    | SelfLeave of string

type SocketStream(socket : Socket) = 
    inherit NetworkStream(socket)
    member  this.Shutdown(how) = this.Socket.Shutdown(how)
    

// Various useful literals for sending and receiving messages
[<Literal>]
let private SPREAD_MAJOR_VERSION = 4uy
[<Literal>]
let private SPREAD_MINOR_VERSION = 0uy
[<Literal>]
let private SPREAD_PATCH_VERSION = 0uy

[<Literal>]
let private MAX_PROC_NAME = 20
[<Literal>]
let private MAX_PRIVATE_NAME = 10
[<Literal>]
let private MAX_GROUP_NAME = 32


[<Literal>]
let private MAX_AUTH_NAME = 30
[<Literal>]
let private MAX_AUTH_METHODS = 3

[<Literal>]
let private HEADER_SIZE = 48
[<Literal>]
let private ENDIAN_TYPE = 0x80000080
[<Literal>]
let private ARCH_ENDIAN = 0x80000080
[<Literal>]
let private OFFSET_TYPE = 0
[<Literal>]
let private OFFSET_SENDER = 4
[<Literal>]
let private OFFSET_NUM_GROUPS = 36
[<Literal>]
let private OFFSET_HINT = 40
[<Literal>]
let private OFFSET_DATA_LEN = 44

// Is our arch the same as the sender?
let private isSameEndian v =
    (v &&& ENDIAN_TYPE) = ARCH_ENDIAN

// Flip the endianess of an int32
let private flipInt32 (buffer : byte array) offset = 
    let b0 = buffer.[offset]
    let b1 = buffer.[offset + 1]
    let b2 = buffer.[offset + 2]
    let b3 = buffer.[offset + 3]
    
    buffer.[offset] <- b3
    buffer.[offset + 1] <- b2
    buffer.[offset + 2] <- b1
    buffer.[offset + 3] <- b0

// Active pattern to match type of the message
let (|Regular|Membership|Transition|SelfLeave|) st = 
    if st &&& MessageFlags.Regular <> MessageFlags.None then
        Regular
    elif st &&& MessageFlags.Membership <> MessageFlags.None then
        if st &&& MessageFlags.Transition = MessageFlags.Transition then Transition
        elif st &&& MessageFlags.RegMembership <> MessageFlags.None then Membership
        elif st &&& MessageFlags.CausedByLeave = MessageFlags.CausedByLeave then SelfLeave
        else failwith "Unknown service type"
    else failwith "Unknown service type"

// Various parsing functions for turning the bits into something comprehensible
let private parseGroup (buffer : byte array ) offset = 
    let idx = Array.IndexOf<byte>(buffer, 0uy, offset, MAX_GROUP_NAME)
    System.Text.Encoding.ASCII.GetString(buffer, offset, idx - offset)

let private parseGroups buffer offset numGroups = 
    let groups = Array.create<string> numGroups null
    for i = 0 to (numGroups - 1) do
        groups.[i] <- parseGroup buffer (offset + i * MAX_GROUP_NAME)
    groups
    
let private parseMessageType hdr offset endianMatch = 
    if not endianMatch then flipInt32 hdr offset
    let bits = BitConverter.ToUInt32(hdr, offset)
    uint16 (((bits  &&& ~~~0x80000080u) >>> 8) &&& 0xFFFFu)

let private parseGroupId buffer offset endianMatch = 
    if not endianMatch then 
        flipInt32 buffer offset
        flipInt32 buffer (offset + 4)
        flipInt32 buffer (offset + 8)
    let pid = BitConverter.ToInt32(buffer, offset)
    let ts = BitConverter.ToInt32(buffer, offset + 4)
    let idx = BitConverter.ToInt32(buffer, offset + 8)
    { ProcessId = pid; Timestamp = ts; Index = idx }

let private parseChangeset buffer offset sameEndian =
    let gid = parseGroupId buffer offset sameEndian
    if  not sameEndian then
        flipInt32 buffer (offset + 12)
        flipInt32 buffer (offset + 16)
    let num = BitConverter.ToInt32(buffer, offset + 12)
    let idx = BitConverter.ToInt32(buffer, offset + 16)
    
    let (sets : string array array) = Array.zeroCreate num
    for i = 0 to (num - 1) do
        let goffset = (offset + 20) + i * (4 + MAX_GROUP_NAME)
        if  not sameEndian then flipInt32 buffer goffset
        let numGroups = BitConverter.ToInt32(buffer, goffset)
        sets.[i] <- parseGroups buffer (goffset + 4) numGroups

    { Gid = gid; IndexOfLocalGroup = idx; GroupMembers = sets }

// Convert the constituent bits into an event
let private eventFromBuffer stype endianMatch hdr groups body  = 
    match stype with
    | Regular ->
        let s = parseGroup hdr OFFSET_SENDER
        let n = BitConverter.ToInt32(hdr, OFFSET_NUM_GROUPS)
        let g = parseGroups groups 0 n
        let mt = parseMessageType hdr OFFSET_HINT endianMatch
        Message { 
            Sender = s; 
            Groups = g; 
            MessageType = mt; 
            Data = body;
            SameEndian = endianMatch }
    | Transition ->
        let s = parseGroup hdr OFFSET_SENDER
        let gid = parseGroupId body 0 endianMatch
        Transition { Group = s; Gid = gid }
    | SelfLeave ->
        SelfLeave (parseGroup hdr OFFSET_SENDER)
    | Membership ->
        let g = parseGroup hdr OFFSET_SENDER
        let n = BitConverter.ToInt32(hdr, OFFSET_NUM_GROUPS)
        let m = parseGroups groups 0 n
        let cs = parseChangeset body 0 endianMatch
        let isSet mf f = mf &&& f = f
        if isSet stype MessageFlags.CausedByJoin then 
            MembershipChange { Group = g; Members = m; Reason = Joined cs.GroupMembers.[0].[0] }
        elif isSet stype MessageFlags.CausedByLeave then 
            MembershipChange { Group = g; Members = m; Reason = Left cs.GroupMembers.[0].[0] }
        elif isSet stype MessageFlags.CausedByDisconnect then 
            MembershipChange { Group = g; Members = m; Reason = Disconnected cs.GroupMembers.[0].[0] }
        elif isSet stype MessageFlags.CausedByNetwork then 
            MembershipChange { Group = g; Members = m; Reason = Partition cs }
        else failwith "Not a valid regular membership message"

// An async loop that reads, decodes and delivers events from the daemon. 
// This implementation is simple and but not as fast as could it be. As such,
// it can be disconnected from the daemon if it fails to process events fast enough.
// I have a more efficient implementation in C#, using pre-posted 
// SocketAsyncEventArgs and may re-implement this along those lines.
let private runRecv (stream : NetworkStream) callback = 
    let flipEndian buffer =
        flipInt32 buffer OFFSET_TYPE
        flipInt32 buffer OFFSET_NUM_GROUPS
        flipInt32 buffer OFFSET_DATA_LEN

    let rec loop =
        async {
            try
                let! hdr = stream.AsyncRead HEADER_SIZE
                let endianMatch = isSameEndian (BitConverter.ToInt32(hdr,0))
                if not endianMatch then flipEndian hdr
                let serviceType = enum<MessageFlags> (BitConverter.ToInt32(hdr, 0))
                let groupLength = BitConverter.ToInt32(hdr, OFFSET_NUM_GROUPS) * MAX_GROUP_NAME
                let dataLength = BitConverter.ToInt32(hdr, OFFSET_DATA_LEN)
                let! groups = stream.AsyncRead groupLength
                let! data = stream.AsyncRead dataLength

                eventFromBuffer serviceType endianMatch hdr groups data |> callback
                do! loop
            with error when error.InnerException = null -> ()
        }
    loop

// An async loop that waits on a mailbox, converts the commands to bytes and transmits them
let private runSend (thisGroup : string) (stream : SocketStream) = 
    let withEndian n = (n &&& ~~~ENDIAN_TYPE) ||| ARCH_ENDIAN
    let groupBytes (g : string) = System.Text.Encoding.ASCII.GetBytes(g.PadRight(MAX_GROUP_NAME, char 0))
    let privateGroup = groupBytes thisGroup
    let serialize serviceType messageType groups (data : byte array) = 
        let ms = new MemoryStream()
        let bw = new BinaryWriter(ms)
        int32 serviceType |> withEndian |> bw.Write
        privateGroup |> bw.Write
        Array.length groups |> bw.Write
        Int32.shift_left (int32 messageType) 8 |> Int32.logand 0xFFFF00 |> withEndian |> bw.Write
        Array.length data |> bw.Write
        Array.iter (groupBytes >> bw.Write) groups
        data |> bw.Write
        ms.ToArray()
        
    let channel = MailboxProcessor.Start(fun inbox ->
        let rec loop = 
            async {
                let! msg = inbox.Receive()
                match msg with 
                | Data d -> 
                    match d.Groups.Length with 
                    | 0 ->
                        do! loop
                    | _ ->
                        let buffer = serialize d.ServiceType d.MessageType d.Groups d.Data
                        do! stream.AsyncWrite(buffer, 0, buffer.Length)              
                        do! loop
                | Join group ->
                    let buffer = serialize MessageFlags.Join 0us [|group|] Array.empty
                    do! stream.AsyncWrite(buffer, 0, buffer.Length)              
                    do! loop
                | Leave group -> 
                    let buffer = serialize MessageFlags.Leave 0us [|group|] Array.empty
                    do! stream.AsyncWrite(buffer, 0, buffer.Length)              
                    do! loop
                | Disconnect ->
                    let buffer = serialize MessageFlags.Kill 0us [|thisGroup|] Array.empty
                    do! stream.AsyncWrite(buffer, 0, buffer.Length)              
                | Kill ->
                    try stream.Shutdown(SocketShutdown.Send) with _ -> ()
            }
        loop)
    channel.Post

// Connect to the daemon, including the initial handshake. No authentication
// with the daemon supported.
let private connect endPoint groupMembership name = 
    let connectAsync (socket : Socket) (endPoint : IPEndPoint) = 
        Async.BuildPrimitive(endPoint, socket.BeginConnect, socket.EndConnect)
    async {
        // Create the socket and establish the connection
        let socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp)
        socket.NoDelay <- true
        do! connectAsync socket endPoint
        
        let stream = new SocketStream(socket)

        // Send the connection message
        let ms = new MemoryStream()
        ms.WriteByte(SPREAD_MAJOR_VERSION)
        ms.WriteByte(SPREAD_MINOR_VERSION)
        ms.WriteByte(SPREAD_PATCH_VERSION)
        ms.WriteByte(match groupMembership with true -> 1uy | false -> 0uy)
        
        match name with 
        | Some (s : string) ->
            let bytes = System.Text.Encoding.ASCII.GetBytes(s)
            let count = Array.length bytes
            ms.WriteByte(Byte.of_int count)
            ms.Write(bytes, 0, count)
        | None -> 
            ms.WriteByte 0uy
        do! stream.AsyncWrite(ms.GetBuffer(), 0, int32 ms.Length)
        
        // Now parse the response
        // Read a single byte back - this is the size of the auth array
        let! authLenBytes = stream.AsyncRead(1)
        let authLen = Byte.to_int authLenBytes.[0] 
        
        // We only support null authentication
        let ourAuth = Array.create (MAX_AUTH_METHODS * MAX_AUTH_NAME) 0uy
        ourAuth.[0] <- 'N'B
        ourAuth.[1] <- 'U'B
        ourAuth.[2] <- 'L'B
        ourAuth.[3] <- 'L'B
        do! stream.AsyncWrite(ourAuth)
        
        // Now check for the server response
        let! connected = stream.AsyncRead(1 + authLen)
        if connected.[authLen] = 0uy then failwith "Server refused connection."
        
        // Read the daemon version
        let! vb = stream.AsyncRead(3)
        let version = (Byte.to_int vb.[0]) * 10000 + (Byte.to_int vb.[1] * 100) + (Byte.to_int vb.[2])
        if version < 30100 then failwith "Old spread version not supported."
                
        return stream
    }
    
let private readPrivateName (stream : NetworkStream) = 
    async {
        let! nameLenBytes = stream.AsyncRead(1)
        let! name = stream.AsyncRead(Byte.to_int nameLenBytes.[0])
        return System.Text.Encoding.ASCII.GetString(name)
    }    

// The top-level function, connects then spawns a reader and writer
let runClient endPoint groupMembership name callback = 
    async {
        try
            let! stream = connect endPoint groupMembership name
            let! groupName = readPrivateName stream
            let sendChannel = runSend groupName stream
            callback (Connection (groupName,sendChannel))
            
            try 
                do! runRecv stream callback
                callback (Disconnection None)
            finally
                sendChannel Kill
        with fault ->
            callback (Disconnection (Some fault))
    } |> Async.Start
    
